from stable_baselines3 import PPO
import gymnasium as gym
from dl_rl.rl_env import TradingEnv
import pandas as pd
import matplotlib.pyplot as plt
from dl_rl.dl_model import TransformerFeatureExtractor
from stable_baselines3.common.policies import ActorCriticPolicy
import os

# 自定义 PPO 策略，使用 Transformer 特征提取器
class CustomPPOPolicy(ActorCriticPolicy):
    def __init__(self, *args, **kwargs):
        super().__init__(
            *args,
            **kwargs,
            features_extractor_class=TransformerFeatureExtractor,
            features_extractor_kwargs=dict(d_model=128, nhead=4, num_layers=2, dim_feedforward=256)
        )

# 创建交易环境
env = TradingEnv(csv_path="./data/20240731/sh/600000/order.csv", symbol="600000", target_quantity=100000000, initial_cash=1000000000, price_log_file="logs/predict/price.csv", is_train=False)

# 加载已经训练好的 PPO 模型，并指定自定义策略
model = PPO.load(
    "D:\proj\python\\bishe\ppo_trading_model.zip",
    env=env,
    custom_objects={"policy_class": CustomPPOPolicy}
)

# 运行一轮交易过程，并记录指标
obs, _ = env.reset()
done = False

records = {
    'step': [],                # 步骤
    'executed_quantity': [],   # 成交数量
    'remaining_quantity': [],  # 剩余数量
    'avg_price': [],           # 平均成交价
    'market_price': [],        # 市场价格
    'cash': [],                # 资金占用
    'completion_rate': [],     # 完成率
    'vwap': []                 # 成交量加权平均价
}

while not done:
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = env.step(action)
    records['step'].append(env.current_step)
    records['executed_quantity'].append(info['executed_quantity'])
    records['remaining_quantity'].append(info['remaining_quantity'])
    records['avg_price'].append(info['avg_price'])
    records['market_price'].append(info['market_price'])
    records['cash'].append(info['current_cash'])
    records['completion_rate'].append(info['completion_rate'])
    records['vwap'].append(info['vwap'])
    done = terminated or truncated

# 转换为 DataFrame 方便绘图
df = pd.DataFrame(records)

# 创建保存图表的目录
os.makedirs("figures", exist_ok=True)

# --- 可视化部分 ---

# 完成率
plt.figure(figsize=(8,4))
plt.plot(df['step'], df['completion_rate'], marker='o', label='Completion Rate')
plt.xlabel("Step")
plt.ylabel("Completion Rate")
plt.title("Order Completion Rate")
plt.grid(True)
plt.legend()
plt.savefig("figures/completion_rate.png")
plt.close()

# 成交均价 vs 市场价格
plt.figure(figsize=(8,4))
plt.plot(df['step'], df['avg_price'], label='Average Executed Price')
plt.plot(df['step'], df['market_price'], label='Market Price', linestyle='--')
plt.xlabel("Step")
plt.ylabel("Price")
plt.title("Average Executed Price vs Market Price")
plt.grid(True)
plt.legend()
plt.savefig("figures/avg_vs_market_price.png")
plt.close()

# 累积成交量
plt.figure(figsize=(8,4))
plt.plot(df['step'], df['executed_quantity'], marker='o')
plt.xlabel("Step")
plt.ylabel("Executed Quantity")
plt.title("Cumulative Executed Quantity")
plt.grid(True)
plt.savefig("figures/executed_quantity.png")
plt.close()

# VWAP
plt.figure(figsize=(8,4))
plt.plot(df['step'], df['vwap'], marker='o', color='purple')
plt.xlabel("Step")
plt.ylabel("VWAP")
plt.title("Volume Weighted Average Price (VWAP)")
plt.grid(True)
plt.savefig("figures/vwap.png")
plt.close()

# 资金变化
plt.figure(figsize=(8,4))
plt.plot(df['step'], df['cash'], marker='x', color='green')
plt.xlabel("Step")
plt.ylabel("Cash Remaining")
plt.title("Cash Usage Over Time")
plt.grid(True)
plt.savefig("figures/cash.png")
plt.close()

# 剩余目标数量
plt.figure(figsize=(8,4))
plt.plot(df['step'], df['remaining_quantity'], marker='o', color='red')
plt.xlabel("Step")
plt.ylabel("Remaining Quantity")
plt.title("Remaining Quantity Over Time")
plt.grid(True)
plt.savefig("figures/remaining_quantity.png")
plt.close()

print("✅ 图表已保存到 ./figures")
